
#include "eye_tiling.h"
#include "register/op_def_registry.h"


namespace optiling {
constexpr int32_t CORE_NUM = 1;
constexpr int32_t UB_SIZE = 196608;

constexpr int32_t X_INDEX = 0;
constexpr int32_t NUM_ROW_ATTR_INDEX = 0;
constexpr int32_t NUM_COL_ATTR_INDEX = 1;
constexpr int32_t BATCH_SHAPE_INDEX = 2;
constexpr int32_t DTYPE_INDEX = 3;
constexpr int32_t BYTE_REPEAT = 256;
// tiling mode
constexpr static int64_t TILING_MODE_SMALL_MATRIX = 0;
constexpr static int64_t TILING_MODE_MID_MATRIX = 1;
constexpr static int64_t TILING_MODE_LARGE_MATRIX = 2;

constexpr int32_t BLOCK_BYTES_SIZE = 32;

constexpr int32_t BUF_COUNT_CACHED = (1) * 1 + 2;
constexpr int32_t BUF_COUNT_NOT_CACHED = (1) * 2 + 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    EyeTilingData tiling;
    uint32_t totalCoreNum = CORE_NUM;
    uint64_t ubSizePlatForm = UB_SIZE;

    
    uint32_t tilingCoreNum = totalCoreNum;
    const gert::Shape xShape = context->GetInputShape(X_INDEX)->GetStorageShape();
    const ge::DataType dataType = context->GetInputDesc(X_INDEX)->GetDataType();
    int32_t dTypeSize = sizeof(dataType);
    int32_t numElemPerBlock = BLOCK_BYTES_SIZE / dTypeSize;
    tiling.set_numElemPerBlock(numElemPerBlock);

    int64_t totalDataCount = xShape.GetShapeSize();
    auto attrs = context->GetAttrs();
    if (attrs == nullptr) {
        return ge::GRAPH_FAILED;
    }
    
    const int64_t* numRowPtr = attrs->GetInt(NUM_ROW_ATTR_INDEX);
    int64_t numRowVal = *numRowPtr;
    tiling.set_numRows(numRowVal);
    
    const int64_t* numColPtr = attrs->GetInt(NUM_COL_ATTR_INDEX);
    int64_t numColVal = *numColPtr;
    const gert::TypedContinuousVector<int64_t>* batchPtr = attrs->GetListInt(BATCH_SHAPE_INDEX);
    int64_t batchNum = *batchPtr->GetData();
    if (batchNum < 1) {
        batchNum = 1;
    }
    tiling.set_batchNum(batchNum);
    const int64_t* dtypePtr = attrs->GetInt(DTYPE_INDEX);
    int64_t dtypeAttr = *dtypePtr;

    if (numColVal <= 0) {
        numColVal = numRowVal;
    }
    
    tiling.set_numColumns(numColVal);

    auto numOnesPerBatch = std::min(numRowVal, numColVal);
    int64_t task_per_core = 0;
    int64_t task_tail = 0;
    int64_t tiling_mode = 0;
    if (numOnesPerBatch * numColVal < numElemPerBlock) { 
        context->SetTilingKey(101);
        task_per_core = 0;
        task_tail = 1;
        tiling_mode = TILING_MODE_SMALL_MATRIX;
    } else if (numColVal < numElemPerBlock){
        context->SetTilingKey(201);
        task_per_core = batchNum / tilingCoreNum;
        task_tail = batchNum % tilingCoreNum;
        tiling_mode = TILING_MODE_MID_MATRIX;
    }

    tiling.set_taskPerCore(task_per_core);

    tiling.set_taskTail(task_tail);
    tiling.set_tilingMode(tiling_mode);

    tiling.set_tilingCoreNum(tilingCoreNum);

    uint64_t ubMaxProcCount = ubSizePlatForm / BLOCK_BYTES_SIZE * numElemPerBlock;


    tiling.set_ubMaxProcCount(ubMaxProcCount);
    context->SetBlockDim(totalCoreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());


    // 以下代码是ascendebug必须的
    size_t* workspaces = context->GetWorkspaceSizes(1);
    workspaces[0] = 16 * 1024 * 1024;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Eye : public OpDef {
public:
    explicit Eye(const char* name) : OpDef(name)
    {
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("num_rows").Int();
        this->Attr("num_columns").AttrType(OPTIONAL).Int(0);
        this->Attr("batch_shape").AttrType(OPTIONAL).ListInt({1});
        this->Attr("dtype").AttrType(OPTIONAL).Int(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Eye);
}
